import random
import math

def distance(a, b):
    return math.sqrt(sum((a[i] - b[i])**2 for i in range(len(a))))

def assign_cluster(x, centers):
    min_dist = float('inf')
    label = -1
    for i, c in enumerate(centers):
        d = distance(x, c)
        if d < min_dist:
            min_dist = d
            label = i
    return label


def compute_center(points):
    dim = len(points[0])
    center = [0] * dim
    for p in points:
        for i in range(dim):
            center[i] += p[i]
    return [x / len(points) for x in center]


def Kmeans(data, k, epsilon=1e-4, iteration=100):
    n = len(data)

    centers = random.sample(data, k)

    for it in range(iteration):
        print(f"Iteration {it+1} ...")

        clusters = [[] for _ in range(k)]
        for x in data:
            label = assign_cluster(x, centers)
            clusters[label].append(x)

        new_centers = []
        for i in range(k):
            if clusters[i]:  
                new_centers.append(compute_center(clusters[i]))
            else:
                new_centers.append(random.choice(data))

        shift = sum(distance(centers[i], new_centers[i]) for i in range(k))
        print(f"Center shift = {shift}")

        if shift < epsilon:
            print("收敛完成！")
            break

        centers = new_centers

    return centers, clusters

if __name__ == "__main__":
    data = [
        [1, 2], [1, 3], [2, 2],
        [8, 9], [9, 8], [8, 8],
        [0, 1], [9, 9]
    ]

    centers, clusters = Kmeans(data, k=2)
    print("\n最终聚类中心：", centers)
    print("聚类结果：", clusters)
